package middleware

import (
	"bytes"
	"io"
	"log/slog"
	"time"

	"github.com/duke-git/lancet/v2/cryptor"
	"github.com/duke-git/lancet/v2/random"
	"github.com/gin-gonic/gin"
)

func RequestLogMiddleware(logger *slog.Logger) gin.HandlerFunc {
	return func(ctx *gin.Context) {
		// The configuration is initialized once per request
		uuid, err := random.UUIdV4()
		if err != nil {
			return
		}
		trace := cryptor.Md5String(uuid)
		ctx.Set("trace", trace)

		// Create a new logger with the context values
		logger = logger.With(
			slog.String("trace", trace),
			slog.String("request_method", ctx.Request.Method),
			slog.String("request_url", ctx.Request.URL.String()),
		)

		if ctx.Request.Body != nil {
			bodyBytes, _ := ctx.GetRawData()
			ctx.Request.Body = io.NopCloser(bytes.NewBuffer(bodyBytes)) // 关键点
			logger = logger.With(slog.String("request_params", string(bodyBytes)))
		}

		// Store the logger in the context
		ctx.Set("logger", logger)
		logger.Info("Request")
		ctx.Next()
	}
}

func ResponseLogMiddleware(logger *slog.Logger) gin.HandlerFunc {
	return func(ctx *gin.Context) {
		blw := &bodyLogWriter{body: bytes.NewBufferString(""), ResponseWriter: ctx.Writer}
		ctx.Writer = blw
		startTime := time.Now()
		ctx.Next()
		duration := time.Since(startTime).String()

		// Retrieve the logger from context if it exists
		if ctxLogger, exists := ctx.Get("logger"); exists {
			if l, ok := ctxLogger.(*slog.Logger); ok {
				logger = l
			}
		}

		logger.Info("Response",
			slog.String("response_body", blw.body.String()),
			slog.String("time", duration),
		)
	}
}

type bodyLogWriter struct {
	gin.ResponseWriter
	body *bytes.Buffer
}

func (w bodyLogWriter) Write(b []byte) (int, error) {
	w.body.Write(b)
	return w.ResponseWriter.Write(b)
}
